#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-i', '--input', help='Path to the input file', default='tc.720x360x30.h0.nc')
parser.add_argument('-o', '--output', help='Path to the output figure')
parser.add_argument('-t', '--time-step', help='Time step to plot', type=int, default=-1)
parser.add_argument('-l', '--loop-plot', help='Loop the plot', action='store_true')
args = parser.parse_args()

fig = plt.figure(figsize=(6, 8))
proj = ccrs.PlateCarree(central_longitude=180)

zlev = np.linspace(0, 16, 150) * units.km

levels = np.linspace(0, 80, 33)
colors = ListedColormap([
	[0.00, 0.11, 0.98],
	[0.00, 0.20, 0.97],
	[0.00, 0.35, 0.98],
	[0.00, 0.52, 0.98],
	[0.00, 0.70, 0.98],
	[0.00, 0.91, 0.99],
	[0.00, 0.98, 0.90],
	[0.00, 0.91, 0.72],
	[0.00, 0.85, 0.55],
	[0.00, 0.79, 0.38],
	[0.00, 0.73, 0.25],
	[0.00, 0.67, 0.15],
	[0.15, 0.72, 0.16],
	[0.32, 0.78, 0.17],
	[0.49, 0.84, 0.18],
	[0.67, 0.90, 0.20],
	[0.87, 0.95, 0.21],
	[1.00, 0.96, 0.22],
	[1.00, 0.88, 0.20],
	[1.00, 0.80, 0.19],
	[1.00, 0.73, 0.18],
	[1.00, 0.66, 0.17],
	[1.00, 0.59, 0.15],
	[1.00, 0.46, 0.13],
	[1.00, 0.34, 0.12],
	[1.00, 0.21, 0.11],
	[1.00, 0.04, 0.10],
	[1.00, 0.00, 0.09],
	[0.94, 0.00, 0.11],
	[0.80, 0.00, 0.15],
	[0.67, 0.00, 0.22],
	[0.54, 0.00, 0.30],
	[0.42, 0.00, 0.38],
	[0.31, 0.02, 0.48],
])

def plot(data_path, fig):
	f = addfile(data_path).isel(time=args.time_step)

	time = parse_datetime(f.time)
	fig.suptitle(f'Day {time.strftime("%d %H:%M")}', fontsize=8)

	z = f.gz * units('m**2 / s**2') / g

	u100m = vinterp_z(z, f.u, np.array([100]) * units.m)
	v100m = vinterp_z(z, f.v, np.array([100]) * units.m)

	wsp100m = np.sqrt(u100m**2 + v100m**2)

	ax = fig.add_subplot(2, 1, 1, projection=proj)
	ax.set_extent([140, 220, 0, 40])
	plot_contour_map(ax, wsp100m[0,:,:], cmap=colors, levels=levels, left_string=f'Wind speed on 100m height (max: {wsp100m.max():.2f})', with_contour=False, with_cbar=False)

	ax = fig.add_subplot(2, 1, 2)
	wsp = vinterp_z(z, np.sqrt(f.u**2 + f.v**2), zlev)
	j, i = np.unravel_index(np.argmin(f.phs.data), f.phs.data.shape)
	plot_contour_lon(ax, wsp.isel(lat=j, lon=slice(i-20, i+20)), cmap=colors, levels=levels, with_contour=False, cbar_orient='horizontal')

if args.loop_plot:
	loop_plot(args.input, fig, plot)
else:
	plot(args.input, fig)
	plt.tight_layout()
	plt.show()
